#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
调度系统连接器
负责与远程调度系统进行通信，包括Worker注册、心跳、任务拉取和状态汇报
"""

import time
import json
import logging
import requests
import threading
import psutil
from typing import Dict, Any, List, Optional, Callable
from datetime import datetime
from dataclasses import dataclass
from enum import Enum

try:
    # 尝试相对导入（当作为模块使用时）
    from .config_manager import get_config
except ImportError:
    # 直接导入（当直接运行时）
    from config_manager import get_config


class TaskStatus(Enum):
    """任务状态枚举"""
    AVAILABLE = "AVAILABLE"
    CLAIMED = "CLAIMED"
    PROCESSING = "PROCESSING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"


@dataclass
class RemoteTask:
    """远程任务数据结构"""
    id: str
    taskType: str
    taskParams: Optional[Dict[str, Any]] = None  # 新字段，可能不在轮询响应中
    priority: int = 0
    createdTime: Optional[str] = None
    deadline: Optional[str] = None

    # 保持向后兼容性的属性
    @property
    def cozeDraftData(self) -> Optional[Dict[str, Any]]:
        """向后兼容的属性访问器"""
        return self.taskParams


class SchedulerConnector:
    """调度系统连接器"""

    def __init__(self, config_manager=None):
        """初始化连接器"""
        self.config = config_manager or get_config()
        self.logger = logging.getLogger(__name__)

        # 获取Worker配置
        self.worker_config = self.config.get_worker_config()

        # 验证必要配置
        if not self.worker_config.get('enabled', False):
            raise ValueError("Worker未启用，请在配置中启用Worker功能")

        required_fields = ['workerName', 'api_base']
        for field in required_fields:
            if not self.worker_config.get(field):
                raise ValueError(f"Worker配置缺少必要字段: {field}")

        # 基础配置
        self.worker_id = None  # 将由注册接口返回
        self.worker_name = self.worker_config['workerName']
        self.api_base = self.worker_config['api_base'].rstrip('/')
        self.max_concurrent_tasks = self.worker_config.get('maxConcurrentTasks', 4)
        self.capabilities = self.worker_config.get('capabilities', ["VIDEO_RENDER", "IMAGE_PROCESS"])
        self.version = self.worker_config.get('version', "1.0.0")

        # 网络配置
        self.connect_timeout = self.worker_config.get('connect_timeout', 30)
        self.max_retries = self.worker_config.get('max_retries', 3)
        self.retry_delay = self.worker_config.get('retry_delay', 5)

        # 心跳和任务轮询配置（将由注册接口更新）
        self.heartbeat_interval = self.worker_config.get('heartbeat_interval', 30)
        self.task_poll_interval = self.worker_config.get('task_poll_interval', 10)

        # 状态管理
        self.is_registered = False
        self.is_running = False
        self.current_task_count = 0

        # 线程管理
        self.heartbeat_thread = None
        self.task_poll_thread = None
        self.stop_event = threading.Event()

        # 任务回调
        self.task_received_callback: Optional[Callable[[RemoteTask], None]] = None

        # 会话管理
        self.session = requests.Session()
        self.session.timeout = self.connect_timeout

        self.logger.info(f"调度系统连接器初始化完成 - Worker ID: {self.worker_id}")

    def set_task_received_callback(self, callback: Callable[[RemoteTask], None]):
        """设置任务接收回调函数"""
        self.task_received_callback = callback

    def register_worker(self) -> bool:
        """注册Worker到调度系统"""
        try:
            # 获取系统信息
            hostname = self.worker_config.get('hostname', 'unknown')
            ip_address = self.worker_config.get('ipAddress', '127.0.0.1')

            # 构建注册数据（按照文档格式）
            registration_data = {
                "workerName": self.worker_name,
                "ipAddress": ip_address,
                "hostname": hostname,
                "maxConcurrentTasks": self.max_concurrent_tasks,
                "capabilities": self.capabilities,
                "version": self.version,
                "systemInfo": json.dumps({
                    "cpu": f"{psutil.cpu_count()}核",
                    "memory": f"{psutil.virtual_memory().total // (1024**3)}GB"
                })
            }

            self.logger.info(f"正在注册Worker: {self.worker_name}")
            self.logger.debug(f"注册数据: {registration_data}")

            # 发送注册请求
            url = f"{self.api_base}/register"
            response = self._make_request('POST', url, json=registration_data)

            if response and response.status_code == 200:
                # 解析注册响应
                response_data = response.json()
                if response_data.get("code") == 200:
                    data = response_data.get("data", {})
                    # 从响应中获取重要配置
                    self.worker_id = data.get("workerId")
                    if "pollInterval" in data:
                        self.task_poll_interval = data["pollInterval"] / 1000  # 转换为秒
                    if "heartbeatInterval" in data:
                        self.heartbeat_interval = data["heartbeatInterval"] / 1000  # 转换为秒

                    self.is_registered = True
                    self.logger.info(f"Worker注册成功: {self.worker_id}")
                    self.logger.info(f"服务器配置 - 轮询间隔: {self.task_poll_interval}s, 心跳间隔: {self.heartbeat_interval}s")
                    return True
                else:
                    error_msg = f"Worker注册失败: {response_data.get('message', '未知错误')}"
                    self.logger.error(error_msg)
                    return False
            else:
                error_msg = f"Worker注册失败: HTTP {response.status_code if response else 'None'}"
                if response:
                    try:
                        error_detail = response.text
                        error_msg += f", {error_detail}"
                    except:
                        pass
                self.logger.error(error_msg)
                return False

        except Exception as e:
            self.logger.error(f"Worker注册异常: {e}")
            return False

    def start_heartbeat(self):
        """启动心跳线程"""
        if self.heartbeat_thread and self.heartbeat_thread.is_alive():
            self.logger.warning("心跳线程已在运行")
            return

        self.stop_event.clear()
        self.heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
        self.heartbeat_thread.start()
        self.logger.info("心跳线程已启动")

    def start_task_polling(self):
        """启动任务轮询线程"""
        if self.task_poll_thread and self.task_poll_thread.is_alive():
            self.logger.warning("任务轮询线程已在运行")
            return

        self.stop_event.clear()
        self.task_poll_thread = threading.Thread(target=self._task_poll_loop, daemon=True)
        self.task_poll_thread.start()
        self.logger.info("任务轮询线程已启动")

    def stop(self):
        """停止所有线程"""
        self.logger.info("正在停止调度系统连接器...")
        self.is_running = False
        self.stop_event.set()

        # 等待线程结束
        if self.heartbeat_thread and self.heartbeat_thread.is_alive():
            self.heartbeat_thread.join(timeout=5)

        if self.task_poll_thread and self.task_poll_thread.is_alive():
            self.task_poll_thread.join(timeout=5)

        # 关闭会话
        self.session.close()

        self.logger.info("调度系统连接器已停止")

    def _heartbeat_loop(self):
        """心跳循环"""
        self.logger.info("心跳循环已启动")

        while not self.stop_event.is_set():
            try:
                if self.is_registered:
                    success = self._send_heartbeat()
                    if not success:
                        self.logger.warning("心跳发送失败，可能需要重新注册")
                        self.is_registered = False

                # 等待下一次心跳
                self.stop_event.wait(self.heartbeat_interval)

            except Exception as e:
                self.logger.error(f"心跳循环异常: {e}")
                self.stop_event.wait(self.retry_delay)

        self.logger.info("心跳循环已停止")

    def _task_poll_loop(self):
        """任务轮询循环"""
        self.logger.info("任务轮询循环已启动")

        while not self.stop_event.is_set():
            try:
                if self.is_registered and self.current_task_count < self.max_concurrent_tasks:
                    tasks = self._poll_tasks()

                    for task in tasks:
                        if self.current_task_count >= self.max_concurrent_tasks:
                            break

                        # 任务已自动分配，直接处理
                        self.logger.info(f"接收到已分配的任务: {task.id}")

                        # 通知任务处理器
                        if self.task_received_callback:
                            try:
                                self.task_received_callback(task)
                                self.current_task_count += 1
                            except Exception as e:
                                self.logger.error(f"任务回调处理异常: {e}")
                                # 任务处理失败，报告失败
                                self._report_task_failed(task.id, f"任务回调处理异常: {str(e)}")

                # 等待下一次轮询
                self.stop_event.wait(self.task_poll_interval)

            except Exception as e:
                self.logger.error(f"任务轮询循环异常: {e}")
                self.stop_event.wait(self.retry_delay)

        self.logger.info("任务轮询循环已停止")

    def _send_heartbeat(self) -> bool:
        """发送心跳"""
        try:
            # 获取系统信息
            system_info = self._get_system_info()

            # 确定worker状态
            status = "BUSY" if self.current_task_count > 0 else "ONLINE"

            heartbeat_data = {
                "workerId": self.worker_id,
                "status": status,
                "currentTasks": self.current_task_count,
                "systemInfo": json.dumps(system_info)  # 按文档要求，应该是JSON字符串
            }

            url = f"{self.api_base}/heartbeat"
            response = self._make_request('POST', url, json=heartbeat_data)

            if response and response.status_code == 200:
                self.logger.debug(f"心跳发送成功: {self.worker_id} (状态: {status})")
                return True
            else:
                self.logger.warning(f"心跳发送失败: HTTP {response.status_code if response else 'None'}")
                return False

        except Exception as e:
            self.logger.error(f"发送心跳异常: {e}")
            return False

    def _poll_tasks(self) -> List[RemoteTask]:
        """轮询可用任务"""
        try:
            url = f"{self.api_base}/tasks/poll"
            params = {"workerId": self.worker_id}

            response = self._make_request('GET', url, params=params)

            if response and response.status_code == 200:
                response_data = response.json()
                if response_data.get("code") == 200:
                    task_data = response_data.get("data")

                    # 检查是否有任务（data为null表示无任务）
                    if task_data is None:
                        self.logger.debug("暂无可用任务")
                        return []

                    # 解析任务数据
                    # taskParams现在直接包含在轮询响应中，可能是JSON字符串格式
                    task_params_raw = task_data.get("taskParams")
                    task_params = None

                    if task_params_raw:
                        try:
                            # 如果taskParams是JSON字符串，解析为字典
                            if isinstance(task_params_raw, str):
                                import json
                                task_params = json.loads(task_params_raw)
                            else:
                                task_params = task_params_raw
                        except Exception as e:
                            self.logger.warning(f"解析taskParams失败: {e}, 原始数据: {task_params_raw}")
                            task_params = None

                    task = RemoteTask(
                        id=task_data["id"],
                        taskType=task_data["taskType"],
                        taskParams=task_params,
                        priority=task_data.get("priority", 0),
                        createdTime=task_data.get("createdTime"),
                        deadline=task_data.get("deadline")
                    )

                    self.logger.debug(f"轮询到1个已分配任务: {task.id}")
                    return [task]
                else:
                    self.logger.warning(f"轮询任务失败: {response_data.get('message', '未知错误')}")
                    return []
            else:
                self.logger.warning(f"轮询任务失败: HTTP {response.status_code if response else 'None'}")
                return []

        except Exception as e:
            self.logger.error(f"轮询任务异常: {e}")
            return []

    def report_task_start(self, task_id: str) -> bool:
        """报告任务开始执行"""
        try:
            url = f"{self.api_base}/tasks/{task_id}/start"
            params = {"workerId": self.worker_id}

            response = self._make_request('POST', url, params=params)

            if response and response.status_code == 200:
                self.logger.debug(f"任务开始报告成功: {task_id}")
                return True
            else:
                self.logger.warning(f"任务开始报告失败: {task_id}")
                return False

        except Exception as e:
            self.logger.error(f"任务开始报告异常: {task_id}, {e}")
            return False

    def report_task_processing(self, task_id: str, progress: int = 0) -> bool:
        """报告任务处理中（兼容性保留）"""
        # 如果进度为0，调用开始接口
        if progress == 0:
            return self.report_task_start(task_id)
        else:
            return self.report_task_progress(task_id, progress)

    def report_task_progress(self, task_id: str, progress: int, status_message: str = None) -> bool:
        """报告任务进度"""
        try:
            url = f"{self.api_base}/tasks/{task_id}/progress"
            progress_data = {
                "workerId": self.worker_id,
                "progress": progress
            }
            if status_message:
                progress_data["statusMessage"] = status_message

            response = self._make_request('PUT', url, json=progress_data)

            if response and response.status_code == 200:
                self.logger.debug(f"任务进度更新成功: {task_id}, 进度: {progress}%")
                return True
            else:
                self.logger.warning(f"任务进度更新失败: {task_id}")
                return False

        except Exception as e:
            self.logger.error(f"任务进度更新异常: {task_id}, {e}")
            return False

    def report_task_completed(self, task_id: str, result: Dict[str, Any], oss_video_url: str = None) -> bool:
        """报告任务完成"""
        try:
            url = f"{self.api_base}/tasks/{task_id}/complete"
            complete_data = {
                "workerId": self.worker_id,
                "ossVideoUrl": oss_video_url
            }

            # 添加视频时长和文件大小信息（如果有的话）
            if result:
                if "videoDuration" in result:
                    complete_data["videoDuration"] = result["videoDuration"]
                if "fileSize" in result:
                    complete_data["fileSize"] = result["fileSize"]

            response = self._make_request('POST', url, json=complete_data)

            if response and response.status_code == 200:
                self.logger.info(f"任务完成报告成功: {task_id}")
                self.current_task_count = max(0, self.current_task_count - 1)
                return True
            else:
                self.logger.error(f"任务完成报告失败: {task_id}")
                return False

        except Exception as e:
            self.logger.error(f"任务完成报告异常: {task_id}, {e}")
            return False

    def _report_task_failed(self, task_id: str, error_message: str, error_code: str = None) -> bool:
        """报告任务失败"""
        try:
            url = f"{self.api_base}/tasks/{task_id}/fail"
            fail_data = {
                "workerId": self.worker_id,
                "errorMessage": error_message
            }
            if error_code:
                fail_data["errorCode"] = error_code

            response = self._make_request('POST', url, json=fail_data)

            if response and response.status_code == 200:
                self.logger.info(f"任务失败报告成功: {task_id}")
                self.current_task_count = max(0, self.current_task_count - 1)
                return True
            else:
                self.logger.error(f"任务失败报告失败: {task_id}")
                return False

        except Exception as e:
            self.logger.error(f"任务失败报告异常: {task_id}, {e}")
            return False

    def report_task_failed(self, task_id: str, error_message: str) -> bool:
        """公开方法：报告任务失败"""
        return self._report_task_failed(task_id, error_message)

    def _get_system_info(self) -> Dict[str, Any]:
        """获取系统信息"""
        try:
            # 获取CPU和内存使用率
            cpu_usage = psutil.cpu_percent(interval=1)
            memory = psutil.virtual_memory()
            memory_usage = memory.percent

            return {
                "cpuUsage": int(cpu_usage),
                "memoryUsage": int(memory_usage)
            }
        except Exception as e:
            self.logger.warning(f"获取系统信息失败: {e}")
            return {
                "cpuUsage": 0,
                "memoryUsage": 0
            }

    def _make_request(self, method: str, url: str, **kwargs) -> Optional[requests.Response]:
        """发送HTTP请求（带重试机制）"""
        for attempt in range(self.max_retries):
            try:
                response = self.session.request(method, url, **kwargs)
                response.raise_for_status()
                return response

            except requests.exceptions.Timeout:
                self.logger.warning(f"请求超时: {method} {url} (尝试 {attempt + 1}/{self.max_retries})")
            except requests.exceptions.ConnectionError:
                self.logger.warning(f"连接错误: {method} {url} (尝试 {attempt + 1}/{self.max_retries})")
            except requests.exceptions.RequestException as e:
                self.logger.warning(f"请求异常: {method} {url}, {e} (尝试 {attempt + 1}/{self.max_retries})")
                return None  # 对于4xx, 5xx错误，不重试
            except Exception as e:
                self.logger.error(f"意外异常: {method} {url}, {e}")
                return None

            if attempt < self.max_retries - 1:
                time.sleep(self.retry_delay)

        self.logger.error(f"请求最终失败: {method} {url}")
        return None

    def get_status(self) -> Dict[str, Any]:
        """获取连接器状态"""
        return {
            "worker_id": self.worker_id,
            "is_registered": self.is_registered,
            "is_running": self.is_running,
            "current_task_count": self.current_task_count,
            "max_concurrent_tasks": self.max_concurrent_tasks,
            "heartbeat_running": self.heartbeat_thread and self.heartbeat_thread.is_alive(),
            "task_polling_running": self.task_poll_thread and self.task_poll_thread.is_alive(),
            "api_base": self.api_base
        }


if __name__ == "__main__":
    # 测试代码
    logging.basicConfig(level=logging.INFO)

    def test_task_callback(task: RemoteTask):
        print(f"收到任务: {task.id}, 类型: {task.taskType}")
        # 这里可以添加任务处理逻辑

    try:
        connector = SchedulerConnector()
        connector.set_task_received_callback(test_task_callback)

        if connector.register_worker():
            connector.is_running = True
            connector.start_heartbeat()
            connector.start_task_polling()

            print("Worker已启动，按Ctrl+C停止...")
            while connector.is_running:
                time.sleep(1)
        else:
            print("Worker注册失败")

    except KeyboardInterrupt:
        print("\n正在停止Worker...")
    except Exception as e:
        print(f"启动失败: {e}")
    finally:
        if 'connector' in locals():
            connector.stop()
        print("Worker已停止")